import ast

import torch

from Transformer import Transformer
from config import *
from vocab import tokenizer, Vocab, mapping


def load_vocab(path):
    with open(path, "r", encoding="utf8") as f:
        rel = ast.literal_eval(f.readline())  # 字符串列表转回列表
    return rel


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if __name__ == "__main__":
    # 加载词汇表同时构造词汇表类
    vocab_en = Vocab(load_vocab("./vocab/vocab_en.txt"), init=False)
    vocab_ch = Vocab(load_vocab("./vocab/vocab_ch.txt"), init=False)

    # 载入模型
    transformer = Transformer()
    transformer.load_state_dict(torch.load("./models/transformer_e2c_T4最新.pkl", map_location=torch.device(device)))
    transformer.eval()
    # 构造编码器的输入
    en_input = ["Sorry to be late."]
    en_token_sentence_list = tokenizer(en_input)
    print("编码器的输入：", en_token_sentence_list)
    vector_en_list = []
    for sentence in en_token_sentence_list:
        vector_en_list.append(mapping(sentence, vocab_en))
    print("编码器输入word2idx：", vector_en_list)
    vector_en_list = torch.tensor(vector_en_list, dtype=torch.int64)

    # 构造解码器的输入端GO
    ch_sentence = []
    ch_token_sentence_list = tokenizer(ch_sentence, inference=True)
    print("解码器的输入：", ch_token_sentence_list)
    vector_ch_list = []
    for sentence in ch_token_sentence_list:
        vector_ch_list.append(mapping(sentence, vocab_ch))
    print("解码器输入word2idx：", vector_ch_list)
    vector_ch_list = torch.tensor(vector_ch_list, dtype=torch.int64)

    for _ in range(Max_len - vector_ch_list.shape[1] + 1):
        output_de = transformer(vector_en_list, vector_ch_list)
        if torch.argmax(output_de.data[:, -1, :], dim=-1).item() == EOS:
            vector_ch_list = torch.cat((vector_ch_list, torch.argmax(output_de.data[:, -1, :], dim=-1).unsqueeze(0)),
                                       dim=-1)
            break
        vector_ch_list = torch.cat((vector_ch_list, torch.argmax(output_de.data[:, -1, :], dim=-1).unsqueeze(0)),
                                   dim=-1)
    rel = mapping(vector_ch_list[0], vocab_ch, reverse=True)
    print("经过transformer的输出：", rel)
    print("#" * 100)
    print("输入：", en_input[0])
    print("输出：", end="")
    for word in rel:
        if word == "<SOS>": continue
        print(word, end="")
    print()
